import tensorflow as tf
import os
from PIL import Image


tf.flags.DEFINE_string('dataset_dir', './train', '')

FLAGS = tf.app.flags.FLAGS


def check_image_size():
    word_dir_list = tf.gfile.ListDirectory(FLAGS.dataset_dir)

    height_list = list()
    width_list = list()
    for word_dir in word_dir_list:
        word_list = tf.gfile.ListDirectory(os.path.join(FLAGS.dataset_dir, word_dir))
        for word in word_list:
            word_file = os.path.join(FLAGS.dataset_dir, word_dir, word)
            img = Image.open(word_file)
            height_list.append(img.size[1])
            width_list.append(img.size[0])

    print(len(height_list))
    print(len(width_list))

    height_set = set(height_list)
    width_set = set(width_list)

    print(height_set)
    print(width_set)


def check_image_count():
    word_dir_list = tf.gfile.ListDirectory(FLAGS.dataset_dir)

    word_size = list()
    for word_dir in word_dir_list:
        word_list = tf.gfile.ListDirectory(os.path.join(FLAGS.dataset_dir, word_dir))
        word_size.append(len(word_list))

    word_size_set = set(word_size)

    print(word_size_set)


check_image_count()